# src/model_train/base_trainer.py
import os
import json
import traceback
import shutil
from datetime import datetime
from typing import Dict, Optional, List, Any
from config.ssq_config import logger, SSQ_CONFIG
import trainer.model_train.model_utils as model_utils


class BaseModelTrainer:
    """模型训练基类(仅封装文件保存和版本管理, 不侵入模型核心逻辑)"""
    def __init__(self, label_col: str, data_dict: Dict, model_type: str,load_version: str = None,run_type: str = 'train'):
        self.run_type = run_type
        self.label_col = label_col
        self.data_dict = data_dict
        self.model_type = model_type
        self.train_version = load_version if load_version is not None else SSQ_CONFIG['train_version']   # 支持传入的版本参数
        self.train_process_root = SSQ_CONFIG['train_file']['model_save_dir']
        if self.run_type == 'predict':
            self.train_process_root = SSQ_CONFIG['predict_file']['model_save_dir']
        # /opt/coding/500w/03.ssq_fractal/train_models/vXGB_train_0.2/r27_next$ cat r27_next_process_info.json
        self.model_process_dir = os.path.join(self.train_process_root, self.train_version, self.label_col)
        self.check_version_exists = self._check_version_exists()
        self.train_process_info: Dict[str, Any] = {}  # 存储完整训练信息(包括所有参数)
        self.model: Optional[Any] = None  # 模型实例

        # 初始化或加载过程信息
        if self.check_version_exists:
            self._load_train_process_info()
        else:
            self._init_train_process_info()

    # -------------------------- 版本检查与初始化(仅文件操作) --------------------------
    def _check_version_exists(self) -> bool:
        required_files = [f"{self.label_col}_process_info.json", f"{self.label_col}_feature_columns.json"]
        logger.info(f'required_files: {required_files}')

        # 调用优化后的工具函数，校验文件存在性+有效性
        is_valid = model_utils.check_version_files(self.model_process_dir, required_files)
        
        if is_valid:
            logger.info(f"版本有效，可续训: {self.model_process_dir}")
        else:
            # 若无效，清空目录（可选，避免残留无效文件）
            if os.path.exists(self.model_process_dir):
                shutil.rmtree(self.model_process_dir)
                logger.warning(f"版本无效，已删除残留文件: {self.model_process_dir}")
            # 重新创建空目录
            os.makedirs(self.model_process_dir, exist_ok=True)
            logger.info(f"版本无效，将重新训练: {self.model_process_dir}")
        
        return is_valid

    def _init_train_process_info(self) -> None:
        """初始化过程信息(仅基础字段, 模型参数由子类填充)"""
        self.train_process_info = {
            "version": self.train_version,
            "label_col": self.label_col,
            "model_type": self.model_type,
            "train_time": datetime.now().strftime("%Y-%m-%d_%H:%M:%S"),
            "device": "",  # 由子类填充
            "data_info": {},  # 由子类填充
            "model_params": {},  # 由子类填充完整参数
            "training_info": {"total_iterations": 0, "retrain_history": []},
            "threshold_info": {},
            "evaluation_metrics": {},
            "feature_info": {},
            "status": "init",
            "random_seed": SSQ_CONFIG.get('random_seed',42)
        }

    # -------------------------- 过程信息加载与保存(仅文件操作) --------------------------
    def _load_train_process_info(self) -> bool:
        """加载完整过程信息(包括所有模型参数)"""
        try:
            process_info_path = os.path.join(self.model_process_dir, f"{self.label_col}_process_info.json")
            logger.info(f'_load_train_process_info() process_info_path: {process_info_path}')
            with open(process_info_path, "r", encoding="utf-8") as f:
                self.train_process_info = json.load(f)
            return True
        except Exception as e:
            logger.error(f"加载训练信息失败: {str(e)}\n{traceback.format_exc()}")
            self._init_train_process_info()
            return False

    def _save_train_process_info(self) -> None:
        """保存完整过程信息(包括所有模型参数)"""
        save_path = os.path.join(self.model_process_dir, f"{self.label_col}_process_info.json")
        try:
            process_info_converted = model_utils.convert_numpy_types(self.train_process_info)
            model_utils.save_file_atomic(process_info_converted, save_path)
        except Exception as e:
            logger.error(f"保存训练信息失败: {str(e)}\n{traceback.format_exc()}")

    # -------------------------- 公共文件保存方法(仅文件操作) --------------------------
    def _save_feature_columns(self, feature_cols: List[str]) -> bool:
        """保存特征列(由子类调用, 不干涉数据处理)"""
        try:
            feat_cols_path = os.path.join(self.model_process_dir, f"{self.label_col}_feature_columns.json")
            model_utils.save_file_atomic(feature_cols, feat_cols_path)
            self.train_process_info["feature_info"]["feature_columns"] = feature_cols
            self.train_process_info["feature_info"]["feature_count"] = len(feature_cols)
            return True
        except Exception as e:
            logger.error(f"保存特征列失败: {str(e)}\n{traceback.format_exc()}")
            return False

    def _save_data_statistics(self, data_stats: Dict) -> bool:
        """保存数据统计(由子类传入统计结果, 不干涉计算)"""
        try:
            stats_path = os.path.join(self.model_process_dir, f"{self.label_col}_data_statistics.json")
            model_utils.save_file_atomic(data_stats, stats_path)
            self.train_process_info["data_info"] = data_stats
            return True
        except Exception as e:
            logger.error(f"保存数据统计失败: {str(e)}\n{traceback.format_exc()}")
            return False

    def _save_threshold(self, threshold_info: Dict) -> bool:
        """保存阈值(由子类传入结果)"""
        try:
            threshold_path = os.path.join(self.model_process_dir, f"{self.label_col}_optimal_threshold.json")
            model_utils.save_file_atomic(threshold_info, threshold_path)
            self.train_process_info["threshold_info"] = threshold_info
            return True
        except Exception as e:
            logger.error(f"保存阈值失败: {str(e)}\n{traceback.format_exc()}")
            return False

    # -------------------------- 子类必须实现的核心逻辑(保留原始流程) --------------------------
    def train(self) -> Dict[str, Any]:
        """模型训练核心逻辑(子类必须实现, 保留所有原始参数和步骤)"""
        raise NotImplementedError("子类必须实现训练逻辑")

    @staticmethod
    def print_training_summary(all_results: List[Dict[str, Any]]) -> None:
        raise NotImplementedError("子类必须实现汇总打印")